#include "encryptor.h"

Encryptor::Encryptor(Context *params) : params(params), ringQ(params->ringQ), ringP(params->ringP)
{
    auto *prng = new KeyedPRNG();
    poolQ.resize(1);
    poolQ[0] = ringQ->NewPoly();

    gaussianSampler = new GaussianSampler(prng, ringQ, params->sigma, uint64_t(6 * params->sigma));
    ternarySampler = new TernarySampler(prng, ringQ, 0.5);
    uniformSampler = new UniformSampler(prng, ringQ);
}

PkEncryptor::PkEncryptor(Context *params, PublicKey *key) : Encryptor(params)
{
    if (key->value[0].Q->GetDegree() != params->N() || key->value[1].Q->GetDegree() != params->N()) {
        throw std::runtime_error("cannot newEncryptor: pk ring degree does not match params ring degree");
    }
    if (params->PCount() <= 0) {
        throw std::runtime_error("cannot newEncryptor: pk ring degree should not be negative");
    }
    pk = key;
}
/*
 * @brief 使用公钥pk对明文进行公钥加密
 *
 * @param pt 指向输入的 `Plaintext` 类型对象的指针，包含要进行公钥加密的明文数据。
 * @param ct 指向输出的 `Ciphertext` 类型对象的指针，包含公钥加密后的密文数据。
 */
void PkEncryptor::Encrypt(Plaintext *pt, Ciphertext *ct)
{
    int levelQ = std::min(pt->Level(), ct->Level());

    // We sample a R-WLE instance (encryption of zero) over the extended ring (ciphertext ring + special Prime)

    bool ciphertextNTT = ct->value[0]->isNTT;
    ternarySampler->ReadLvl(levelQ, poolQ[0]);
    // (#Q + #P) NTT
    ringQ->NTTLvl(levelQ, poolQ[0], poolQ[0]);

    ringQ->MulCoeffsBarrettLvl(levelQ, poolQ[0], pk->value[0].Q, ct->value[0]); // -1-1
    ringQ->MulCoeffsBarrettLvl(levelQ, poolQ[0], pk->value[1].Q, ct->value[1]);
    ringQ->InvNTTLvl(levelQ, ct->value[0], ct->value[0]);
    ringQ->InvNTTLvl(levelQ, ct->value[1], ct->value[1]);

    gaussianSampler->ReadLvl(levelQ, poolQ[0]);
    ringQ->AddLvl(levelQ, ct->value[0], poolQ[0], ct->value[0]);
    gaussianSampler->ReadLvl(levelQ, poolQ[0]);
    ringQ->AddLvl(levelQ, ct->value[1], poolQ[0], ct->value[1]);

    if (ciphertextNTT) {
        if (!pt->value->isNTT) {
            ringQ->AddLvl(levelQ, ct->value[0], pt->value, ct->value[0]);
        }

        // 2*#Q NTT
        ringQ->NTTLvl(levelQ, ct->value[0], ct->value[0]);
        ringQ->NTTLvl(levelQ, ct->value[1], ct->value[1]);

        if (pt->value->isNTT) {
            // ct0 = (u*pk0 + e0)/P + m
            ringQ->AddLvl(levelQ, ct->value[0], pt->value, ct->value[0]);
        }
    } else {
        if (!pt->value->isNTT) {
            ringQ->AddLvl(levelQ, ct->value[0], pt->value, ct->value[0]);
        } else {
            ringQ->InvNTTLvl(levelQ, pt->value, poolQ[0]);
            ringQ->AddLvl(levelQ, ct->value[0], poolQ[0], ct->value[0]);
        }
    }

    ct->value[1]->isNTT = ct->value[0]->isNTT;
    ct->value[0]->Coeffs.resize(levelQ + 1);
    ct->value[1]->Coeffs.resize(levelQ + 1);
}

/*
 * @brief 对明文数据进行加密并返回新的密文对象。
 *
 * 该函数接受一个明文对象 `pt`，并使用当前加密参数生成新的密文对象。通过调用 `Encrypt`
 * 函数执行加密操作。返回的密文对象是一个新的 `Ciphertext` 实例。
 *
 * @param pt 输入的明文对象，类型为 `Plaintext`，将被加密为密文。
 *
 * @return 返回加密后的 `Ciphertext` 对象，类型为 `Ciphertext*`，此对象包含加密结果。
 */
Ciphertext *PkEncryptor::EncryptNew(Plaintext *pt)
{
    auto *ct = new Ciphertext(params, 1, params->MaxLevel());
    Encrypt(pt, ct);
    return ct;
}

SkEncryptor::SkEncryptor(Context *params, SecretKey *key) : Encryptor(params)
{
    if (key->value.Q->GetDegree() != params->N()) {
        throw std::runtime_error("cannot newEncryptor: sk ring degree does not match params ring degree");
    }
    sk = key;
}
/*
 * @brief 使用私钥sk对明文进行私钥加密
 *
 * @param pt 指向输入的 `Plaintext` 类型对象的指针，包含要进行私钥加密的明文数据。
 * @param ct 指向输出的 `Ciphertext` 类型对象的指针，包含私钥加密后的密文数据。
 */
void SkEncryptor::Encrypt(Plaintext *pt, Ciphertext *ct)
{
    int levelQ = std::min(pt->Level(), ct->Level());

    bool ciphertextNTT = ct->value[0]->isNTT;
    ringQ->MulCoeffsBarrettLvl(levelQ, ct->value[1], sk->value.Q, ct->value[0]); // -1-1
    ringQ->NegLvl(levelQ, ct->value[0], ct->value[0]);

    if (ciphertextNTT) {
        gaussianSampler->ReadLvl(levelQ, poolQ[0]);
        if (pt->value->isNTT) {
            ringQ->NTTLvl(levelQ, poolQ[0], poolQ[0]);
            ringQ->AddLvl(levelQ, ct->value[0], poolQ[0], ct->value[0]);
            ringQ->AddLvl(levelQ, ct->value[0], pt->value, ct->value[0]);
        } else {
            ringQ->AddLvl(levelQ, poolQ[0], pt->value, poolQ[0]);
            ringQ->NTTLvl(levelQ, poolQ[0], poolQ[0]);
            ringQ->AddLvl(levelQ, ct->value[0], poolQ[0], ct->value[0]);
        }
        ct->value[0]->isNTT = true;
        ct->value[1]->isNTT = true;
    } else {
        if (pt->value->isNTT) {
            ringQ->AddLvl(levelQ, ct->value[0], pt->value, ct->value[0]);
            ringQ->InvNTTLvl(levelQ, ct->value[0], ct->value[0]);
        } else {
            ringQ->InvNTTLvl(levelQ, ct->value[0], ct->value[0]);
            ringQ->AddLvl(levelQ, ct->value[0], pt->value, ct->value[0]);
        }

        gaussianSampler->ReadAndAddLvl(levelQ, ct->value[0]);
        ringQ->InvNTTLvl(levelQ, ct->value[1], ct->value[1]);

        ct->value[0]->isNTT = false;
        ct->value[1]->isNTT = false;
    }

    ct->value[0]->Coeffs.resize(levelQ + 1);
    ct->value[1]->Coeffs.resize(levelQ + 1);
}
/*
 * @brief 使用私钥对明文数据进行加密并返回新的密文对象。
 *
 * 该函数接受一个明文对象 `pt`，并使用当前加密参数生成新的密文对象。调用 `Encrypt`
 * 函数执行加密操作。返回的密文对象是一个新的 `Ciphertext` 实例。
 *
 * @param pt 输入的明文对象，类型为 `Plaintext`，将被加密为密文。
 *
 * @return 返回加密后的 `Ciphertext` 对象，类型为 `Ciphertext*`，此对象包含加密结果。
 */
Ciphertext *SkEncryptor::EncryptNew(Plaintext *pt)
{
    auto *ct = new Ciphertext(params, 1, params->MaxLevel());
    Encrypt(pt, ct);
    return ct;
}
